use std::{collections::HashMap, num::NonZeroU16};

#[derive(Default)]
struct TrieNode {
    children: HashMap<char, u32>,
    table_index: Option<NonZeroU16>,
}

/// 用于简繁体转换的字典树
pub struct Trie {
    nodes: Vec<TrieNode>,
    table: Vec<String>,
}
impl Trie {
    fn new() -> Self {
        Self {
            nodes: vec![TrieNode::default()],
            table: Vec::new(),
        }
    }

    fn insert(&mut self, s: &str, t: impl Into<String>) {
        assert!(self.table.len() <= u16::MAX as usize);
        let table_index = self.table.len() as u16;
        self.table.push(t.into());

        let mut node_index = 0;
        for c in s.chars() {
            node_index = match self.nodes[node_index].children.get(&c) {
                Some(&n) => n,
                None => {
                    let n = self.nodes.len() as u32;
                    self.nodes.push(TrieNode::default());
                    self.nodes[node_index].children.insert(c, n);
                    n
                }
            } as usize;
        }
        self.nodes[node_index].table_index = NonZeroU16::new(table_index);
    }

    /// 使用最大正向匹配翻译字符串
    pub fn translate(&self, s: &str) -> String {
        let chars: Vec<char> = s.chars().collect();

        let mut i = 0;
        let mut result = String::new();
        while i < chars.len() {
            let mut node_index = 0;
            let mut last_match = None;
            for (len, &c) in chars[i..].iter().enumerate() {
                if let Some(&n) = self.nodes[node_index].children.get(&c) {
                    node_index = n as usize;
                    if let Some(t) = self.nodes[node_index].table_index {
                        last_match = Some((t.get(), len + 1));
                    }
                } else {
                    break;
                }
            }
            if let Some((t, len)) = last_match {
                result += &self.table[t as usize];
                i += len;
            } else {
                result.push(chars[i]);
                i += 1;
            }
        }

        result
    }
}

thread_local! {
    static TRIES: (Trie, Trie) = build_tries();
}

fn build_tries() -> (Trie, Trie) {
    let mut s2t = Trie::new();
    let mut t2s = Trie::new();

    let data = include_str!("dict.txt");
    for line in data.lines() {
        let (simplified, traditional) = line.split_once(' ').unwrap();
        s2t.insert(simplified, traditional);
        t2s.insert(traditional, simplified);
    }

    (s2t, t2s)
}

pub fn converter(input: &str, tp: &str) -> String {
    TRIES.with(|(s2t, t2s)| match tp {
        "s2t" => s2t.translate(input),
        "t2s" => t2s.translate(input),
        _ => input.to_owned(),
    })
}
